import pandas as pd
import pickle
from utils import *
from model import *
import os
import sys


from optparse import OptionParser

usage = "test.py --test-dataset=TESTING_SET_DIR --model=MODDIR --prediction-file=PREDICTION"
parser= OptionParser(usage)

def parse_argv():
    parser.add_option("-e", "--test-dataset", action="store", metavar="DIR", type="string", dest="data_dir",
                      help="directory for testing set")
    parser.add_option("-m", "--model", action="store", metavar="DIR", type="string", dest="model_dir",
                      help="directory for training set")
    parser.add_option("-p", "--prediction-file", action="store", metavar="DIR", type="string", dest="predict_file",
                      help="directory for training set")
    if len(sys.argv) <3:
        parser.print_help()
    #    sys.exit(0)
    else:
        (options, args) = parser.parse_args()
        return options.data_dir, options.model_dir, options.predict_file


def load_test(input_dir, model_dir):
    testfile = os.listdir(input_dir)[0]
    df = pd.read_csv(input_dir + "/" + testfile)
    dt_from = gen_test_submit(df, outpath=model_dir +"/test1.csv")
    print("predicting from {}{}".format(dt_from, type(dt_from)))
    df = pd.read_csv(model_dir +"/test1.csv")
    fill_and_shift(df, model_dir +"/test2", shift_days=28)
    df = pd.read_csv(model_dir +"/test2.csv")

    return dt_from, df

def run_predicting(df, category_fields, shift_fields, yfields, model_dir):

    for i in range(len(yfields)):


        yfield = yfields[i]

        print("loading model for fields:[{}]".format(yfield))
        with open(model_dir + "/"+yfield + ".model", 'rb') as f:
            model = pickle.load(f)
            print(model)
        with open(model_dir + "/"+yfield + ".ohe", 'rb') as f:
            encoder = pickle.load(f)

        x_fields =  shift_fields

        print("predicting ({}) samples ".format(len(df)))
        x_test = convert_testdf_to_matrix(df, category_fields, x_fields, encoder)

        y_test = df[yfield]

        y_pred = predicting(x_test, model)

        score = model.score(x_test, y_test)
        print("test score:{}".format(score))
        df["pred".format(i)] = y_pred
        print("===========")

        df_submit = df[['date','station', yfield,"pred" ]].sort_values(['station','date'])
        df_submit.to_csv(model_dir + '/submit_{}.csv'.format(yfield), index=False)

def r2_score(y_pred, y_test):
    y_pred = np.array(y_pred)
    y_test = np.array(y_test)
    y_test_mean = y_test.mean()
    numerator = ((y_pred - y_test) * (y_pred - y_test)).sum()
    denominator = ((y_test - y_test_mean) * (y_test - y_test_mean)).sum()
    return 1- (numerator/denominator)

def merge_submit(model_dir, dt_from, predict_file):
    df_submit1 = pd.read_csv(model_dir + "/submit_temperature.csv")
    df_submit2 = pd.read_csv(model_dir + "/submit_humidity.csv")

    df_test1 = df_submit1[(df_submit1['date'] <= dt_from) ]
    df_test2 = df_submit2[(df_submit2['date'] <= dt_from) ]

    test_score1 = r2_score(df_test1['temperature'], df_test1['pred'])
    test_score2 = r2_score(df_test2['humidity'], df_test2['pred'])
    print("#{} temperature r2:{}\thumidity r2:{}".format(len(df_test1), test_score1, test_score2))

    t_pred = df_submit1[df_submit1['date'] > dt_from]
    h_pred = df_submit2[df_submit2['date'] > dt_from]['pred']
    t_pred['humidity'] = h_pred
    t_pred['temperature'] = t_pred['pred']


    t_pred[['date','station','temperature','humidity']].to_csv(predict_file, index=False)

if __name__ == "__main__":
    input_dir, model_dir, predict_file = parse_argv()

    dt_from, df = load_test(input_dir, model_dir)

    yfields = [ 'temperature', 'humidity']

    df = feature_engineering(df, reload=False, for_train= False)


    category_fields = [ 'month', 'city', 'station']
   # shift_fields = ['humidity', 'temperature', 'rain', 'cloud']
    avg_fields = ['humidity', 'temperature', 'cloud','rain']
    numeric_fields = ['days']
    numeric_fields += get_avg_nfields(avg_fields, FD, WS, SP)
   # numeric_fields = ['days'] + ["{}{}".format(sf, d) for sf in shift_fields for d in [1,2]]
    run_predicting(df, category_fields, numeric_fields, yfields[:1], model_dir)

    run_predicting(df, category_fields, numeric_fields, yfields[1:], model_dir)

    merge_submit(model_dir, dt_from, predict_file)